from langchain_community.document_loaders import WebBaseLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter

# List of URLs to load documents from
urls = [
    "https://lilianweng.github.io/posts/2023-06-23-agent",
    # "https://lilianweng.github.io/posts/2023-03-15-prompt-engineering",
    # "https://lilianweng.github.io/posts/2023-10-25-adv-attack-llm",
]
# Load documents from the URLs
docs = [WebBaseLoader(url).load() for url in urls]
docs_list = [item for sublist in docs for item in sublist]

# Initialize a text splitter with specified chunk size and overlap
text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
    chunk_size=250, chunk_overlap=0
)
# Split the documents into chunks
print("*** docs_list:", docs_list)
doc_splits = text_splitter.split_documents(docs_list)


from langchain_community.vectorstores import SKLearnVectorStore
from langchain_mistralai import MistralAIEmbeddings
# Create embeddings for documents and store them in a vector store
print("*** doc_splits:", doc_splits)
vectorstore = SKLearnVectorStore.from_documents(
    ## 4 is proper number, larger will be too many requests. Smaller will throw error
    documents=doc_splits[0:4],
    #todo: remove api key
    embedding=MistralAIEmbeddings(model="mistral-embed", api_key="1lnTB5FzycBDT7xstCVOPLhJCBvatncN"),
)
retriever = vectorstore.as_retriever(k=1)


from langchain_ollama import ChatOllama
from langchain.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
# Define the prompt template for the LLM
prompt = PromptTemplate(
    template="""You are an assistant for question-answering tasks.
    Use the following documents to answer the question.
    If you don't know the answer, just say that you don't know.
    Use three sentences maximum and keep the answer concise:
    Question: {question}
    Documents: {documents}
    Answer:
    """,
    input_variables=["question", "documents"],
)

# Initialize the LLM with Llama 3.1 model
llm = ChatOllama(
    model="llama3.1:8b",
    temperature=0,
)

# Create a chain combining the prompt template and LLM
rag_chain = prompt | llm | StrOutputParser()


# Define the RAG application class
class RAGApplication:
    def __init__(self, retriever, rag_chain):
        self.retriever = retriever
        self.rag_chain = rag_chain
    def run(self, question):
        # Retrieve relevant documents
        documents = self.retriever.invoke(question)
        # Extract content from retrieved documents
        print("*** Retrived docs:",documents)
        doc_texts = "\\n".join([doc.page_content for doc in documents])
        # Get the answer from the language model
        answer = self.rag_chain.invoke({"question": question, "documents": doc_texts})
        return answer
    

if __name__ == "__main__":
    # Initialize the RAG application
    rag_application = RAGApplication(retriever, rag_chain)
    # Example usage
    question = "What is prompt engineering"
    answer = rag_application.run(question)
    print("Question:", question)
    print("Ollama Answer:", answer)
